from openai import OpenAI
import os
from dotenv import load_dotenv

# pip install openai dotenv


def get_api_key():
    """获取API KEY"""
    load_dotenv()
    api_key = os.getenv("DEEPSEEK_API_KEY")
    return api_key


def get_base_url(platform="deepseek"):
    """获取BASE URL"""
    if platform == "deepseek":
        # print("已经获取到了BASE_URL，平台是deepseek")
        return "https://api.deepseek.com"
    elif platform == "openai":
        # print("已经获取到了BASE_URL，平台是openai")
        return "https://api.openai.com"
    else:
        print("不支持的平台")
        return None


def ai_chat(api_key, base_url, prompt, **kwargs):
    """
    调用AI接口进行对话
    :param api_key: API KEY
    :param base_url: BASE URL
    :param prompt: 提问内容
    :param kwargs: 其他参数
    :return: AI回答内容
    """

    client = OpenAI(api_key=api_key, base_url=base_url)

    prompts = f"""
    你是一个人工智能助手，你的任务是回答用户的问题。
    请用中文回答，不要用英文回答。
    用户的问题如下：{prompt}
    """

    response = client.chat.completions.create(
        model=kwargs.get("model", "deepseek-chat"),
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": f"{prompts}"},
        ],
        temperature=kwargs.get("temperature", 1.0),
        max_tokens=kwargs.get("max_tokens", 4096),
        stream=False,
    )

    contents = response.choices[0].message.content
    return contents


if __name__ == "__main__":
    api_key = get_api_key()
    base_url = get_base_url()

    while True:
        prompt = input("请输入您的问题：")
        response_content = ai_chat(
            api_key, base_url, prompt, temperature=0.9, max_tokens=150
        )
        print(f"AI回复：{response_content}")
